'''
Author: Ryan
Email: 409266255@qq.com
Date: 2025-09-29 17:16:57

LastEditors: Ryan
LastEditTime: 2025-09-29 17:52:41
Description: 头部注释

FilePath: /qopilot/qopilot/utils.py
Copyright (c) 2025 by Ryan, All Rights Reserved. 
'''


import os
import logging
from qopilot.config import config

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def standardize_column_names(df, column_config=None):
    """
    更灵活的列重命名和过滤函数

    Args:
        df: 输入的DataFrame
        column_config: 列配置字典，格式为：
            {
                'rename': {'原列名': '新列名'},  # 重命名映射
                'keep': ['列名1', '列名2'],     # 要保留的列（重命名后）
                'drop': ['列名1', '列名2']      # 要删除的列（重命名后）
            }

    Returns:
        DataFrame: 处理后的DataFrame
    """
    if column_config is None:
        return df

    df_result = df.copy()

    # 重命名列
    if 'rename' in column_config and column_config['rename']:
        df_result = df_result.rename(columns=column_config['rename'])

    # 删除列
    if 'drop' in column_config and column_config['drop']:
        df_result = df_result.drop(columns=column_config['drop'], errors='ignore')

    # 保留指定列
    if 'keep' in column_config and column_config['keep']:
        available_columns = [col for col in column_config['keep'] if col in df_result.columns]
        df_result = df_result[available_columns]

    return df_result


def write_to_excel(df, excel_filename):
    # 将df写入Excel文件
    # 确保目标目录存在
    # 使用相对路径，从当前文件位置（qopilot/data/datasource/）到项目根目录的database文件夹
    database_dir = config.DATABASE_DIR
    os.makedirs(database_dir, exist_ok=True)

    excel_path = os.path.join(database_dir, excel_filename)
    df.to_excel(excel_path, index=False, engine='openpyxl')
    logger.info(f"港股数据已导出到Excel文件: {excel_path}")
